﻿/*
* Copyright (c) 2012 Kyle Hayward
* 
* Permission is hereby granted, free of charge, to any person obtaining a copy
* of this software and associated documentation files (the "Software"), to deal
* in the Software without restriction, including without limitation the rights
* to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
* copies of the Software, and to permit persons to whom the Software is
* furnished to do so, subject to the following conditions:
* 
* The above copyright notice and this permission notice shall be included in
* all copies or substantial portions of the Software.
* 
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
* IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
* FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
* AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
* LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
* OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
* THE SOFTWARE.
*/
using System;
using System.Collections.Generic;
using System.ComponentModel;
using System.Data;
using System.Drawing;
using System.Linq;
using System.Text;
using System.IO;
using System.Windows.Forms;
using System.Runtime.InteropServices;

using OptixDotNet;
using OptixDotNet.Math;
using OptixDotNet.Utils;

using OptixDotNetUtils;

namespace PathTracer
{
    struct ParallelogramLight
    {
        public Vector3 corner;
        public Vector3 v1, v2;
        public Vector3 normal;
        public Vector3 emission;
        public bool textured;
    };

    public class PathTracer : OptixForm
    {
        uint mFrame = 0;
        public PathTracer()
        {
            //Ensure that we have a valid Nvidia GPU
            Utils.DetectGPUCompatibility();
            UseSRGB = true;
        }

        protected override void Initialize()
        {
            string rayGenPath = Path.GetFullPath( "../../Assets/Cuda/path_tracer.ptx" );
            string shaderPath = Path.GetFullPath( "../../Assets/Cuda/path_tracer.ptx" );

            //string modelName = "cornell-dragon.obj";
            string modelName = "cornell-box.obj";
            string modelPath = Path.GetFullPath( "../../Assets/Models/" + modelName );
            string accelCache = "pathtracer." + modelName + ".accel";

            /*-----------------------------------------------
             * Create the Optix context
             *-----------------------------------------------*/
            Context = new Context();
            Context.RayTypeCount = 2;
            Context.EntryPointCount = 1;
            Context.EnableAllExceptions = false;

            /*-----------------------------------------------
             * Create the material that will be executed when there is an intersection
             *-----------------------------------------------*/
            Material material = new Material( Context );
            material.Programs[ 0 ] = new SurfaceProgram( Context, RayHitType.Closest, shaderPath, "diffuse" );
            material.Programs[ 1 ] = new SurfaceProgram( Context, RayHitType.Any, shaderPath, "shadow" );

            /*-----------------------------------------------
             * Load the geometry
             *-----------------------------------------------*/
            OptixOBJLoader model = new OptixOBJLoader( modelPath, Context, null, material );
            model.GeoGroup = new GeometryGroup( Context );
            model.ParseNormals = false;
            model.GenerateNormals = false;

            string intersectPath = Path.GetFullPath( "../../Assets/Cuda/triangle_mesh.ptx" );
            model.IntersecitonProgPath = intersectPath;
            model.BoundingBoxProgPath = intersectPath;
            model.IntersecitonProgName = "mesh_intersect";
            model.BoundingBoxProgName = "mesh_bounds";

            model.LoadContent();

            /*-----------------------------------------------
             * Create scene lights
             *-----------------------------------------------*/
            CreateLights();

            /*-----------------------------------------------
             * Create the output buffer
             *-----------------------------------------------*/
            CreateOutputBuffer( Format.Float4 );

            /*-----------------------------------------------
             * Create the ray-generation and exception programs
             *-----------------------------------------------*/
            Program rayGen = new Program( Context, rayGenPath, "pathtrace_camera" );
            Program exception = new Program( Context, rayGenPath, "exception" );
            Program miss = new Program( Context, shaderPath, "miss" );
            miss[ "bg_color" ].Set( 100 / 255.0f, 149 / 255.0f, 237 / 255.0f );

            Context.SetRayGenerationProgram( 0, rayGen );
            Context.SetExceptionProgram( 0, exception );
            Context.SetRayMissProgram( 0, miss );

            /*-----------------------------------------------
             * Finally compile the optix context, and build the accel tree
             *-----------------------------------------------*/
            SetCamera( model.BBox );

            // mark the acceleration data as dirty if we didn't successfully load a cached accel file
            if ( !LoadAccelData( model, modelPath, accelCache ) )
            {
                model.GeoGroup.Acceleration.MarkAsDirty();
            }

            Context[ "top_object" ].Set( model.GeoGroup );
            Context[ "output_buffer" ].Set( OutputBuffer );
            Context[ "scene_epsilon" ].Set( 0.0003f );

            const uint sqrtSamples = 2u;
            const uint maxDepth = 3u;
            const uint rrBeginDepth = 1u;

            Context[ "rr_begin_depth" ].Set( rrBeginDepth );
            Context[ "max_depth" ].Set( maxDepth );
            Context[ "sqrt_num_samples" ].Set( sqrtSamples );
            Context[ "frame_number" ].Set( mFrame );

            Context[ "pathtrace_ray_type" ].Set( 0u );
            Context[ "pathtrace_shadow_ray_type" ].Set( 1u );

            Context[ "bad_color" ].SetFloat3( new Vector3( 1.0f, 0.0f, 0.0f ) );

            Console.Write( "Compiling Optix... " );
            float start = Time.GetTimeInSecs();

            bool accelIsDirty = model.GeoGroup.Acceleration.IsDirty();
            Context.Compile();
            Context.BuildAccelTree();

            //only save accel data if it was dirty before building
            if ( accelIsDirty )
            {
                model.GeoGroup.Acceleration.Save( accelCache );
            }

            Console.WriteLine( "{0:0.00}s", Time.GetTimeInSecs() - start );

            //give aproximate here becuase we'll have somewhere between rrBeginDepth and maxDepth number of iterations per sample.
            int avgIteration = (int)( rrBeginDepth + maxDepth ) / 2;
            RaysTracedPerFrame = Width * Height * 2 * ( (int)sqrtSamples * (int)sqrtSamples ) * avgIteration;
        }

        private void CreateLights()
        {
            ParallelogramLight light = new ParallelogramLight();
            light.corner = new Vector3( 2.8096f, 17.1165f, 2.3659f );
            light.v1 = new Vector3( -5.0f, 0.0f, 0.0f );
            light.v2 = new Vector3( 0.0f, 0.0f, 5.0f );

            light.normal = light.v1.Cross( light.v2 );
            light.normal.Normalize();
            light.emission = new Vector3( 15.0f, 15.0f, 5.0f );

            BufferDesc desc = new BufferDesc() { Width = 1u, Format = Format.User, Type = BufferType.Input, ElemSize = (uint)Marshal.SizeOf( typeof( ParallelogramLight ) ) };
            OptixDotNet.Buffer lightBuffer = new OptixDotNet.Buffer( Context, desc );
            BufferStream stream = lightBuffer.Map();
            stream.Write( light );
            lightBuffer.Unmap();

            Context[ "lights" ].Set( lightBuffer );
        }

        private void SetCamera( BoundingBox box )
        {
            Camera = new Camera();
            Camera.Aspect = (float)Width / (float)Height;
            Camera.Fov = 35.0f;
            Camera.RotationVel = 100.0f;
            Camera.TranslationVel = 100.0f;

            Camera.CenterOnBoundingBox( box, .95f );

            CameraUpdate();
        }

        protected override void RayTrace()
        {
            Context[ "frame_number" ].Set( mFrame++ );
            Context.Launch( 0, (uint)Width, (uint)Height );
        }

        protected override void Display()
        {
            base.Display();
        }

        protected override void CameraUpdate()
        {
            mFrame = 0;

            Vector3 eye = Camera.Position;
            Vector3 right = Camera.Right;
            Vector3 up = Camera.Up;
            Vector3 look = Camera.Look;

            Context[ "eye" ].Set( ref eye );
            Context[ "U" ].Set( ref right );
            Context[ "V" ].Set( ref up );
            Context[ "W" ].Set( ref look );
        }
    }
}
